from torch.nn import CrossEntropyLoss
from omegaconf import ListConfig
from torch import tensor, Tensor


class CrossEntropyLoss(CrossEntropyLoss):
    def __init__(self, weight = None, size_average=None, ignore_index = -100, reduce=None, reduction = "mean", label_smoothing = 0):
        if weight and not isinstance(weight, Tensor):
            if isinstance(weight, int) or isinstance(weight, float):
                weight = [weight]
            weight = tensor(weight) 
        super().__init__(weight, size_average, ignore_index, reduce, reduction, label_smoothing)
    
    def forward(self, input, target):
        if len(target.shape) > 2 and target.shape[1] == 1:
            # remove channel dimension if class labels
            target = target.squeeze(1)
        return super().forward(input, target.long())